/*
 * Copyright 2021 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package androidx.appsearch.platformstorage;

import android.annotation.SuppressLint;
import android.content.Context;
import android.os.Build;

import androidx.annotation.RequiresApi;
import androidx.annotation.RestrictTo;
import androidx.appsearch.app.GenericDocument;
import androidx.appsearch.app.SearchResult;
import androidx.appsearch.app.SearchResults;
import androidx.appsearch.app.SearchSpec;
import androidx.appsearch.exceptions.AppSearchException;
import androidx.appsearch.platformstorage.converter.SearchResultToPlatformConverter;
import androidx.appsearch.platformstorage.util.AppSearchVersionUtil;
import androidx.collection.ArraySet;
import androidx.concurrent.futures.ResolvableFuture;
import androidx.core.util.Preconditions;

import com.google.common.util.concurrent.ListenableFuture;

import org.jspecify.annotations.NonNull;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.Executor;

/**
 * Platform implementation of {@link SearchResults} which proxies to the platform's
 * {@link android.app.appsearch.SearchResults}.
 * @exportToFramework:hide
 */
@RestrictTo(RestrictTo.Scope.LIBRARY_GROUP)
@RequiresApi(Build.VERSION_CODES.S)
class SearchResultsImpl implements SearchResults {
    private final android.app.appsearch.SearchResults mPlatformResults;
    private final SearchSpec mSearchSpec;
    private final Executor mExecutor;
    private final Context mContext;

    SearchResultsImpl(
            android.app.appsearch.@NonNull SearchResults platformResults,
            @NonNull SearchSpec searchSpec,
            @NonNull Executor executor,
            @NonNull Context context) {
        mPlatformResults = Preconditions.checkNotNull(platformResults);
        mSearchSpec = Preconditions.checkNotNull(searchSpec);
        mExecutor = Preconditions.checkNotNull(executor);
        mContext = Preconditions.checkNotNull(context);
    }

    @SuppressLint("WrongConstant")
    @Override
    public @NonNull ListenableFuture<List<SearchResult>> getNextPageAsync() {
        ResolvableFuture<List<SearchResult>> future = ResolvableFuture.create();
        mPlatformResults.getNextPage(mExecutor, result -> {
            if (result.isSuccess()) {
                List<android.app.appsearch.SearchResult> frameworkResults = result.getResultValue();
                List<SearchResult> jetpackResults = new ArrayList<>(frameworkResults.size());
                Map<String, List<String>> projection = mSearchSpec.getProjections();

                boolean isBuildVersionBeforeUDC =
                        Build.VERSION.SDK_INT < Build.VERSION_CODES.UPSIDE_DOWN_CAKE;
                boolean isAppSearchMainlineVersionBeforeUBase =
                        AppSearchVersionUtil.getAppSearchVersionCode(mContext)
                            < AppSearchVersionUtil.APPSEARCH_U_BASE_VERSION_CODE;
                // Projection must be manually applied on versions before Android U.
                // This is a workaround in Jetpack code for a pre-existing projection bug
                // that persists Android S and T.
                boolean manuallyApplyProjection = isBuildVersionBeforeUDC
                        && isAppSearchMainlineVersionBeforeUBase
                        && !projection.isEmpty();

                for (int i = 0; i < frameworkResults.size(); i++) {
                    if (Build.VERSION.SDK_INT == Build.VERSION_CODES.S
                            || Build.VERSION.SDK_INT == Build.VERSION_CODES.S_V2) {
                        // This is a patch for b/197361770, framework-appsearch in Android S will
                        // disable the whole namespace filter if none of given namespaces exist.
                        // And that will result in Icing returns all documents that this query is
                        // able to access.
                        if (i == 0 && !mSearchSpec.getFilterNamespaces().isEmpty()
                                && !mSearchSpec.getFilterNamespaces().contains(
                                frameworkResults.get(i).getGenericDocument().getNamespace())) {
                            // We should never return a document with a namespace that is not
                            // required in the request. And also since the bug will only happen
                            // when the required namespace doesn't exist, we should just return
                            // an empty result when we found the result contains unexpected
                            // namespace.
                            future.set(Collections.emptyList());
                            return;
                        }
                    }
                    SearchResult jetpackResult =
                            SearchResultToPlatformConverter.toJetpackSearchResult(
                                    frameworkResults.get(i));

                    if (manuallyApplyProjection) {
                        SearchResult searchResult =
                                manuallyApplyProjection(jetpackResult, projection);
                        jetpackResults.add(searchResult);
                    } else {
                        jetpackResults.add(jetpackResult);
                    }
                }
                future.set(jetpackResults);
            } else {
                future.setException(
                        // Without the SuppressLint annotation on the method, this line causes a
                        // lint error because getResultCode isn't defined as returning a value from
                        // AppSearchResult.ResultCode
                        new AppSearchException(result.getResultCode(), result.getErrorMessage()));
            }
        });
        return future;
    }

    @Override
    public void close() {
        mPlatformResults.close();
    }

    /**
     * Helper function to manually apply projection.
     *
     * @param jetpackResult {@link SearchResult} from Jetpack.
     * @param projection {@link SearchSpec} projection.
     * @return {@link SearchResult} with manually applied projection.
     */
    @NonNull
    private SearchResult manuallyApplyProjection(
            @NonNull SearchResult jetpackResult,
            @NonNull Map<String, List<String>> projection) {
        GenericDocument jetpackGenericDocument = jetpackResult.getGenericDocument();
        GenericDocument.Builder jetpackGenericDocumentBuilder =
                new GenericDocument.Builder(jetpackGenericDocument);

        // Determine the properties we want to keep for a given schema type.
        Set<String> retainedProperties = null;
        List<String> schemaTypes = projection.get(jetpackGenericDocument.getSchemaType());
        List<String> wildcardSchemaTypes = projection.get(SearchSpec.SCHEMA_TYPE_WILDCARD);

        if (schemaTypes != null) {
            retainedProperties = new ArraySet<>(schemaTypes);
        } else if (wildcardSchemaTypes != null) {
            retainedProperties = new ArraySet<>(wildcardSchemaTypes);
        }

        if (retainedProperties != null) {
            for (String propertyName : jetpackGenericDocument.getPropertyNames()) {
                GenericDocument propertyDocument = jetpackGenericDocument
                        .getPropertyDocument(propertyName);

                // Check if the property is a nested document.
                if (propertyDocument != null) {
                    // Apply projection to the nested documents.
                    applyProjectionToNestedDocument(jetpackGenericDocumentBuilder,
                            propertyDocument, propertyName, retainedProperties);
                } else if (!retainedProperties.contains(propertyName)) {
                    // If a property is not a part of the projection, it will be removed.
                    jetpackGenericDocumentBuilder.clearProperty(propertyName);
                }
            }
        }

        // Build an updated SearchResult that accounts for projection.
        SearchResult.Builder resultBuilder = new SearchResult.Builder(jetpackResult)
                .setGenericDocument(jetpackGenericDocumentBuilder.build());

        // The updated SearchResult should not hold MatchInfos that have propertyPaths
        // corresponding any removed properties.
        List<SearchResult.MatchInfo> matchInfos = jetpackResult.getMatchInfos();
        if (!matchInfos.isEmpty()) {
            resultBuilder.clearMatchInfos();
        }
        for (SearchResult.MatchInfo matchInfo : matchInfos) {
            String propertyPath = matchInfo.getPropertyPath();

            boolean isPropertyInSchemaTypeProjection = schemaTypes != null
                    && schemaTypes.contains(propertyPath);

            boolean isPropertyInWildcardProjection = wildcardSchemaTypes != null
                    && wildcardSchemaTypes.contains(propertyPath);

            if (isPropertyInSchemaTypeProjection || isPropertyInWildcardProjection) {
                SearchResult.MatchInfo clonedMatchInfo =
                        new SearchResult.MatchInfo.Builder(matchInfo).build();
                resultBuilder.addMatchInfo(clonedMatchInfo);
            }
        }

        // Return the updated SearchResult.
        return resultBuilder.build();
    }

    /**
     * Helper function to manually apply projection to a nested document.
     *
     * @param jetpackGenericDocumentBuilder The builder for the parent {@link GenericDocument}.
     * @param nestedDocument The nested {@link GenericDocument} whose properties need projection.
     * @param propertyName The property name of the parent document.
     * @param retainedProperties The set of properties that need to be retained.
     */
    private void applyProjectionToNestedDocument(
            GenericDocument.@NonNull Builder jetpackGenericDocumentBuilder,
            @NonNull GenericDocument nestedDocument,
            @NonNull String propertyName,
            @NonNull Set<String> retainedProperties) {
        Set<String> nestedProperties = new ArraySet<>();

        for (String property : retainedProperties) {
            if (property.startsWith(propertyName + ".")) {
                // The property will be in the format "propertyName.nestedPropertyName"
                // Capture the "nestedPropertyName"
                nestedProperties.add(property.substring(propertyName.length() + 1));
            }
        }

        if (!nestedProperties.isEmpty()) {
            // Build a new nested document with the projected properties.
            GenericDocument.Builder nestedDocumentBuilder = new GenericDocument
                    .Builder(nestedDocument);
            for (String nestedPropertyName : nestedDocument.getPropertyNames()) {
                GenericDocument nestedPropertyDocument = nestedDocument
                        .getPropertyDocument(nestedPropertyName);

                // Check if the property is a nested document.
                if (nestedPropertyDocument != null) {
                    // Recursively apply projection to the nested documents.
                    applyProjectionToNestedDocument(
                            nestedDocumentBuilder,
                            nestedPropertyDocument,
                            nestedPropertyName,
                            nestedProperties);
                } else if (!nestedProperties.contains(nestedPropertyName)) {
                    // If the nested property is not part of the projection, it will be removed.
                    nestedDocumentBuilder.clearProperty(nestedPropertyName);
                }
            }

            // Modify the builder of the parent document.
            jetpackGenericDocumentBuilder.setPropertyDocument(propertyName,
                    nestedDocumentBuilder.build());
        }
    }
}
